package com.by.rpc.registry;

import com.by.log.Logger;
import com.by.rpc.dto.ServiceInstance;
import com.by.rpc.util.KeyUtils;
import com.by.log.LoggerFactory;

import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 服务心跳检测守护线程
 */
public class HeartbeatChecker {
    // 添加自定义日志实例
    private static final Logger log = LoggerFactory.getLogger(HeartbeatChecker.class);
    
    private final ServiceRegistry serviceRegistry;
    private final ServiceInstance localServiceInstance;
    private final ScheduledExecutorService executorService;
    private static final long HEARTBEAT_INTERVAL = 10; // 心跳间隔10秒
    private static final int MAX_FAILURE_COUNT = 3; // 最大连续失败次数

    private final AtomicInteger failureCount = new AtomicInteger(0); // 连续失败计数
    
    public HeartbeatChecker(ServiceRegistry serviceRegistry, ServiceInstance localServiceInstance) {
        this.serviceRegistry = serviceRegistry;
        this.localServiceInstance = localServiceInstance;
        this.executorService = Executors.newSingleThreadScheduledExecutor(r -> {
            Thread t = new Thread(r, "HeartbeatChecker-" + localServiceInstance.getServiceName());
            t.setDaemon(true);
            return t;
        });
    }

    /**
     * 启动心跳检测
     */
    public void start() {
        // 启动定时任务，定期发送心跳
        executorService.scheduleAtFixedRate(this::checkAndSendHeartbeat, 0, HEARTBEAT_INTERVAL, TimeUnit.SECONDS);
        log.info("Heartbeat checker started for service: {}", localServiceInstance.getServiceName());
    }

    /**
     * 停止心跳检测
     */
    public void stop() {
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
        log.info("Heartbeat checker stopped for service: {}", localServiceInstance.getServiceName());
    }

    /**
     * 检查并发送心跳
     */
    private void checkAndSendHeartbeat() {
        try {
            // 检查服务是否还在注册中心
            String serviceKey = KeyUtils.buildServiceKey(localServiceInstance.getServiceName(),
                    localServiceInstance.getVersion(), localServiceInstance.getGroup());
            List<ServiceInstance> instances = serviceRegistry.discover(serviceKey);
            boolean isRegistered = instances.stream().anyMatch(instance -> 
                instance.getHost().equals(localServiceInstance.getHost()) && 
                instance.getPort() == localServiceInstance.getPort());

            if (!isRegistered) {
                // 服务不在注册中心，需要重新注册
                log.warn("Service {} is not registered, re-registering...", localServiceInstance.getServiceName());
                localServiceInstance.setRegisterTime(System.currentTimeMillis());
                serviceRegistry.register(localServiceInstance);
                failureCount.set(0); // 重置失败计数
                return;
            }

            // 发送心跳
            sendHeartbeat();
        } catch (Exception e) {
            log.error("Failed to check and send heartbeat for service: {}", localServiceInstance.getServiceName(), e);
            handleFailure();
        }
    }

    /**
     * 发送心跳
     */
    private void sendHeartbeat() {
        try {
            localServiceInstance.setRegisterTime(System.currentTimeMillis());
            serviceRegistry.updateHeartbeat(localServiceInstance);
            log.debug("Heartbeat sent for service: {}", localServiceInstance.getServiceName());
            
            // 心跳成功，重置失败计数
            failureCount.set(0);
        } catch (Exception e) {
            log.error("Failed to send heartbeat for service: {}", localServiceInstance.getServiceName(), e);
            handleFailure();
        }
    }

    /**
     * 处理心跳失败
     */
    private void handleFailure() {
        int currentFailureCount = failureCount.incrementAndGet();
        log.warn("Heartbeat failed for service: {}, consecutive failure count: {}",
                localServiceInstance.getServiceName(), currentFailureCount);
        
        // 如果连续失败次数超过阈值，从注册中心移除服务
        if (currentFailureCount >= MAX_FAILURE_COUNT) {
            try {
                log.error("Consecutive heartbeat failures exceeded threshold for service: {}, deregistering...",
                         localServiceInstance.getServiceName());
                serviceRegistry.deregister(localServiceInstance);
                failureCount.set(0); // 重置失败计数
            } catch (Exception e) {
                log.error("Failed to deregister service: {}", localServiceInstance.getServiceName(), e);
            }
        }
    }
}